from functools import partial

import numpy as np
import torch
from datasets import load_dataset
from torch.utils.data import DataLoader
from transformers import default_data_collator


class DataManager:
    """
    数据管理器
    """

    def __init__(self, config, tokenizer):
        self.config = config
        self.tokenizer = tokenizer

    def get_data_loder(self):
        # 模板转换器定义
        dataset = load_dataset('text', data_files={'train': self.config.train_path,
                                                   'dev': self.config.valid_path})
        new_func = partial(self.convert_example,
                           tokenizer=self.tokenizer,
                           p_embedding_num=self.config.p_embedding_num,
                           max_seq_len=self.config.max_seq_len,
                           max_label_len=self.config.max_label_len)

        dataset = dataset.map(new_func, batched=True)

        # self.convert_example(dataset, tokenizer=self.tokenizer,
        #                      hard_template=self.hard_template,
        #                      max_seq_len=self.config.max_seq_len,
        #                      max_label_len=self.config.max_label_len)

        train_dataset = dataset["train"]
        dev_dataset = dataset["dev"]
        train_dataloader = DataLoader(train_dataset,
                                      shuffle=True,
                                      collate_fn=default_data_collator,
                                      batch_size=self.config.batch_size)
        dev_dataloader = DataLoader(dev_dataset,
                                    collate_fn=default_data_collator,
                                    batch_size=self.config.batch_size)
        return train_dataloader, dev_dataloader

    @staticmethod
    def convert_example(examples: dict,
                        tokenizer,
                        max_seq_len: int,
                        max_label_len: int,
                        p_embedding_num: int = 6,
                        train_mode: bool = True,
                        return_tensor: bool = False) -> dict:
        """
        将样本数据转换为模型接收的输入数据。

        Args:
            examples (dict): 训练数据样本, e.g. -> {
                                                    "text": [
                                                                '手机	这个手机也太卡了。',
                                                                '体育	世界杯为何迟迟不见宣传',
                                                                ...
                                                    ]
                                                }
            tokenizer :
            max_seq_len (int): 句子的最大长度，若没有达到最大长度，则padding为最大长度
            max_label_len (int): 最大label长度，若没有达到最大长度，则padding为最大长度
            p_embedding_num (int): p-tuning token 的个数
            train_mode (bool): 训练阶段 or 推理阶段。
            return_tensor (bool): 是否返回tensor类型，如不是，则返回numpy类型。

        Returns:
            dict (str: np.array) -> tokenized_output = {
                                'input_ids': [[101, 3928, ...], [101, 4395, ...]],
                                'token_type_ids': [[0, 0, ...], [0, 0, ...]],
                                'mask_positions': [[5, 6, ...], [3, 4, ...]],
                                'mask_labels': [[183, 234], [298, 322], ...]
                            }
        """
        tokenized_output = {
            'input_ids': [],
            'attention_mask': [],
            # 记录label的位置（即MASK Token的位置）
            'mask_positions': [],
            # 记录MASK Token的原始值（即Label值）
            'mask_labels': []
        }

        for i, example in enumerate(examples['text']):
            try:
                # 将 prompt token(s) 插在 [CLS] 之后
                start_mask_position = 1
                if train_mode:
                    label, content = example.strip().split('\t', 1)
                else:
                    content = example.strip()
                encoded_inputs = tokenizer(
                    text=content,
                    truncation=True,
                    max_length=max_seq_len,
                    padding='max_length')
            except:
                continue
            input_ids = encoded_inputs['input_ids']
            # print(f'encoded_inputs-->{encoded_inputs}')
            # print(f'input_ids-->{input_ids}')

            # 1.生成 MASK Tokens, 和label长度一致
            mask_tokens = ['[MASK]'] * max_label_len
            # print(f'mask_tokens-->{mask_tokens}')

            mask_ids = tokenizer.convert_tokens_to_ids(mask_tokens)  # token 转 id
            # print(f'mask_ids-->{mask_ids}')

            # 2.构建 prompt token(s)
            p_tokens = ["[unused{}]".format(i + 1) for i in range(p_embedding_num)]
            # print(f'p_tokens-->{p_tokens}')
            # token 转 id
            p_tokens_ids = tokenizer.convert_tokens_to_ids(p_tokens)
            # print(f'p_tokens_ids-->{p_tokens_ids}')
            tmp_input_ids = input_ids[:-1]
            # print(f'tmp_input_ids-->{len(tmp_input_ids)}')

            # 根据最大长度-p_token长度-label长度，裁剪content的长度
            tmp_input_ids = tmp_input_ids[:max_seq_len - max_label_len - p_embedding_num - 1]
            # print(f'tmp_input_ids1-->{tmp_input_ids}')
            # print(len(tmp_input_ids))

            # 3.插入 MASK -> [CLS][MASK][MASK]世界杯...[SEP]
            tmp_input_ids = tmp_input_ids[:start_mask_position] + mask_ids + tmp_input_ids[start_mask_position:]
            # print(f'tmp_input_ids2--{tmp_input_ids}')

            # 补上[SEP]
            input_ids = tmp_input_ids + [input_ids[-1]]
            # print(f'input_ids11-->{input_ids}')

            # 4.插入 prompt -> [unused1][unused2]...[CLS][MASK]...[SEP]
            input_ids = p_tokens_ids + input_ids
            # print(f'input_ids22-->{input_ids}')
            # print(f'input_ids33-->{len(input_ids)}')

            # 将 Mask Tokens 的位置记录下来
            mask_positions = [p_embedding_num + start_mask_position + i for i in range(max_label_len)]

            tokenized_output['input_ids'].append(input_ids)

            # 如果输入需要token_type_ids，可以进行添加
            # 兼容不需要 token_type_id 的模型, e.g. Roberta-Base
            if 'token_type_ids' in encoded_inputs:
                tmp = encoded_inputs['token_type_ids']
                if 'token_type_ids' not in tokenized_output:
                    tokenized_output['token_type_ids'] = [tmp]
                else:
                    tokenized_output['token_type_ids'].append(tmp)
            tokenized_output['attention_mask'].append(encoded_inputs['attention_mask'])
            tokenized_output['mask_positions'].append(mask_positions)

            if train_mode:
                mask_labels = tokenizer(text=label)  # label token 转 id
                mask_labels = mask_labels['input_ids'][1:-1]  # 丢掉[CLS]和[SEP]
                mask_labels = mask_labels[:max_label_len]
                # 将 label 补到最长
                mask_labels += [tokenizer.pad_token_id] * (max_label_len - len(mask_labels))
                tokenized_output['mask_labels'].append(mask_labels)

        for k, v in tokenized_output.items():
            if return_tensor:
                tokenized_output[k] = torch.LongTensor(v)
            else:
                tokenized_output[k] = np.array(v)

        return tokenized_output
